from ast import Dict
from datetime import datetime
import re
from typing import List

from loguru import logger
from docx import Document as DocxDocument
from langchain_core.documents import Document


def deduplicate_documents_generator(documents: List[DocxDocument], revert_table=False):
	"""根据文档内容去重的生成器"""
	seen = set()
	for doc in documents:
		if doc.page_content not in seen:
			seen.add(doc.page_content)
			yield doc


def create_documents_from_pdf(pdf_metadata: Dict, file_obj, processed_attachments: List[Dict]) \
        -> List[Document]:
    """从PDF元数据创建文档
    {
        "_source": {
          "text": "...",
          "metadata": {
            "create_time": "2025-08-01T19:03:03.552981",
            "update_time": "2025-08-01T19:03:03.552981",
            "file_id": "bf33c5da-ed06-11ef-802c-00163e0c52c3",
            "file_name": "gb-xxx-xxx.pdf",
            "gb_num": "gb-xxx-xxx",
            "applicable_scope": "...",
            "images": [
                {
                    "id": "002e250bcd0bd520.jpg",
                    'type': 'image',
                    'image_id': "image_001",
                    'content': "图片链接: https://web-api.textin.com/ocr_image/external/002e250bcd0bd520.jpg",
                    "image_url": "https://web-api.textin.com/ocr_image/external/002e250bcd0bd520.jpg",
                    "local_save_path": "/Users/wzq/Desktop/portAI/portai-server/document_loader/test_data/gb-file/002e250bcd0bd520.jpg"
                }
            ],
            "tables": [
                {
                    "id": "table_001",
                    "type": "table",
                    "table_id": "table_001",
                    "content": "表格内容",
                }
            ]
          }
        }
    }
    """
    documents = []

    try:
        # 兼容两种结构：优先 main_text，否则回退到 chunks
        pages = pdf_metadata.get("main_text")
        if not pages:
            chunks = pdf_metadata.get("chunks", [])
            pages = [{"content": c.get("content", "")} for c in chunks]

        for text_obj in pages:
            page_content = text_obj.get("content", "")
            # 同时兼容占位标记和 <img> 标签，且支持 image_*/img_* 两种前缀
            match_image_results = []
            match_image_results.extend(re.findall(r'<--\s*((?:image|img)_\w+)\s*-->', page_content))
            match_image_results.extend(re.findall(r'<img[^>]+src=["\']((?:image|img)_\w+)["\']', page_content))
            match_table_results = re.findall(r'<-- (table_\w+) -->', page_content)

            # 当前页匹配到的图片与表格
            images = [i for i in processed_attachments if i.get("image_id", "") in match_image_results]
            tables = [i for i in processed_attachments if i.get("table_id", "") in match_table_results]

            # 将属于表格的图片归集到对应table的images字段（依据table_id）
            if tables:
                # 所有图片中，找出带table_id的（image_in_table或图片带有table_id）
                images_with_table = [img for img in processed_attachments if img.get("type") in ("image", "image_in_table") and img.get("table_id")]
                table_id_to_images = {}
                for img in images_with_table:
                    tid = img.get("table_id")
                    table_id_to_images.setdefault(tid, []).append(img)

                enriched_tables = []
                for t in tables:
                    tid = t.get("table_id")
                    if tid and tid in table_id_to_images:
                        t = {**t, "images": table_id_to_images.get(tid, [])}
                    enriched_tables.append(t)
                tables = enriched_tables

            document = Document(
                page_content=page_content,
                metadata={
                    "file_id": str(getattr(file_obj, 'id', '') or ''),
                    "file_name": getattr(file_obj, 'name', '') or '',
                    "images": images,
                    "tables": tables,
                    "create_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                    "update_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                }
            )
            documents.append(document)

        return documents
    except Exception as e:
        logger.error(f"Failed to create documents from PDF: {str(e)}")
        raise ValueError(f"Document creation error: {str(e)}", {
            "file_name": getattr(file_obj, 'name', ''),
            "file_id": getattr(file_obj, 'id', ''),
            "attachment_count": len(processed_attachments)
        })

